
#include <time.h>

#include <set>
#include <algorithm>

#include "book.h"
#include "words.h"


struct result_t {

    // [n_doc] -> {n_word -> weight}
    typedef std::vector< std::map<size_t,double> > tfidf_t;

    // [n_doc] -> title
    typedef std::vector<std::string> titles_t;

    // {n_word -> [(n_doc,weight)]}
    typedef std::map< size_t, std::vector< std::pair<size_t,double> > > index_t;

    struct book_t {
        size_t n;
        double w;
        std::set<size_t> words;

        book_t() : n(0), w(0) {}

        book_t(bool, book_t& b) {
            n = b.n;
            w = b.w;
            words.swap(b.words);
        }
    };

    // [n_doc] -> [doc]
    //typedef std::vector< std::vector<book_t> > closest_t;

    tfidf_t tfidf;
    //closest_t closest;
    titles_t titles;
    index_t index;
};

void read_tfidf(const std::string& file, result_t& res) {

    result_t::tfidf_t& tfidf = res.tfidf;
    result_t::titles_t& titles = res.titles;

    std::cout << "Reading tfidf... " << file << std::endl;

    Source fin(file);

    while (1) {

        tfidf.push_back(std::map<size_t,double>());

        Book book;

        if (!read_digest(book, tfidf.back(), fin)) {
            tfidf.pop_back();
            break;
        }

        std::string title = book.title;
        title += ":";

        bool first = true;
        for (const auto& z : book.author) {

            if (first) {
                first = false;
            } else {
                title += ",";
            }

            if (z[0].size() > 0) {
                title += " ";
                title += z[0];
            }

            if (z[1].size() > 0) {
                title += " ";
                title += z[1];
            }

            if (z[2].size() > 0) {
                title += " ";
                title += z[2];
            }
        }

        titles.push_back(title);
    }
}

void print_title(size_t n, result_t::titles_t& titles) {

    std::cout << ":: " << titles.at(n) << std::endl;
}

void print_closest(result_t::titles_t& titles, result_t::book_t& book) {

    std::cout << "  " << titles.at(book.n) << " " << book.w << "    [";

    for (size_t ww : book.words) {
        std::cout << " " << words().get(ww);
    }

    std::cout << " ]" << std::endl;
}

void find_closest(result_t& res) {

    result_t::tfidf_t& tfidf = res.tfidf;
    result_t::index_t& index = res.index;
    result_t::titles_t& titles = res.titles;

    std::cout << "Building index... " << std::endl;

    size_t n = 0;

    for (const auto& b : tfidf) {

        for (const auto& z : b) {

            index[z.first].push_back(std::make_pair(n, z.second));
        }

        ++n;
    }

    std::cout << "Finding pairwise similarity..." << std::endl;
    
    n = 0;

    for (const auto& b : tfidf) {

        std::map<size_t,result_t::book_t> other;

        for (const auto& z : b) {

            for (const auto& y : index[z.first]) {
                
                result_t::book_t& book = other[y.first];

                book.w += (1.0 + z.second) * (1.0 + y.second);
                book.words.insert(z.first);
            }
        }

        std::multimap<double,result_t::book_t*> resorted;

        for (auto& z : other) {
            
            result_t::book_t& book = z.second;
            book.n = z.first;

            resorted.insert(resorted.end(), std::make_pair(book.w, &(book)));
        }

        print_title(n, titles);

        static const size_t N = 10;

        result_t::book_t* prev_book = nullptr;


        auto ri = resorted.rbegin();
        size_t n0 = 0;
        while (ri != resorted.rend() && n0 < N) {

            result_t::book_t& book = *(ri->second);

            bool ok = (book.w >= 2 && book.w < 100);

            if (prev_book != nullptr) {

                if (book.words == prev_book->words ||
                    titles.at(prev_book->n) == titles.at(book.n)) {

                    ok = false;
                }
            }

            if (ok) {

                print_closest(titles, book);

                ++n0;

                if (n0 >= N)
                    break;

                prev_book = &book;
            }

            ++ri;
        }

        ++n;
    }
}

int main(int argc, char** argv) {
    
    try {

        result_t res;

        for (int i = 1; i < argc; ++i) {

            std::string books(argv[i]);

            read_tfidf(books + ".tfidf", res);
        }

        find_closest(res);

    } catch (std::exception& e) {
        std::cerr << "ERROR: " << e.what() << std::endl;
        return 1;

    } catch (...) {
        std::cerr << "UNKNOWN ERROR" << std::endl;
        return 1;
    }

    return 0;
}

